﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using RaytracerLib.MathLib;
using System.Drawing;
using YAXLib;
using TomanuExtensions;
using RaytracerLib.Renderers;
using RaytracerLib.MathLib.RootFinders;

namespace RaytracerLib.Samplers
{
    public class PoissonSampler : NonUniformSampler
    {
        [YAXNode]
        public int Repeats = 20;

        [YAXNode]
        public double Density = 1;

        private double m_cell_size;
        private double m_radius;
        private int m_grid_width;
        private int m_grid_height;
        private Random m_random;
        private int m_width;
        private int m_height;
        private int[,] m_grid;
        List<int> m_actives = new List<int>();
        List<Vector2> m_samples = new List<Vector2>();

        private void PrepareSamples()
        {
            m_width = Film.Width;
            m_height = Film.Height;

            m_radius = 1.0 / Density;
            m_cell_size = m_radius / Math.Sqrt(2);
            m_grid_width = (int)(m_width / m_cell_size) + 1;
            m_grid_height = (int)(m_height / m_cell_size) + 1;

            m_grid = new int[m_grid_width, m_grid_height];
            m_grid.Fill(-1);

            AddFirstPoint();

            while (m_actives.Count != 0)
            {
                int index = m_random.Next(m_actives.Count);

                bool found = false;

                for (int i = 0; i < Repeats; i++)
                    found |= AddNextPoint(m_actives[index]);

                if (!found)
                    m_actives.RemoveAt(index);
            }
        }

        private Vector2 GenerateRandomAround(int a_index)
        {
            double radius = (m_radius + m_radius * m_random.NextDouble());
            double angle = 2 * MathExtensions.PI * m_random.NextDouble();
            Vector2 center = m_samples[a_index];
            return new Vector2(center.X + radius * Math.Sin(angle), center.Y + radius * Math.Cos(angle));
        }

        private bool AddNextPoint(int a_index)
        {
            Vector2 q = GenerateRandomAround(a_index);

            if ((q.X < 0) || (q.X >= m_width) || (q.Y < 0) || (q.Y >= m_height))
                return false;

            int index_x = GetXIndex(q);
            int index_y = GetYIndex(q);

            for (int x = Math.Max(0, index_x - 2); x < Math.Min(m_grid_width, index_x + 3); x++)
            {
                for (int y = Math.Max(0, index_y - 2); y < Math.Min(m_grid_height, index_y + 3); y++)
                {
                    int index = m_grid[x, y];

                    if (index != -1)
                    {
                        if ((q - m_samples[index]).Length < m_radius)
                            return false;
                    }
                }
            }

            m_grid[index_x, index_y] = m_samples.Count;
            m_actives.Add(m_samples.Count);
            m_samples.Add(q);

            return true;
        }

        private int GetXIndex(Vector2 a_point)
        {
            return (int)(a_point.X / m_cell_size);
        }

        private int GetYIndex(Vector2 a_point)
        {
            return (int)(a_point.Y / m_cell_size);
        }

        private void AddFirstPoint()
        {
            Vector2 p = new Vector2(m_width * m_random.NextDouble(), m_height * m_random.NextDouble());

            m_grid[GetXIndex(p), GetYIndex(p)] = m_samples.Count;
            m_actives.Add(m_samples.Count);
            m_samples.Add(p);
        }

        internal override void RenderStart(RenderStartPhase a_phase)
        {
            base.RenderStart(a_phase);

            if (a_phase == RenderStartPhase.PrepareObjectToRender)
            {
                m_random = new System.Random(Film.Scene.RenderOptions.RandomSeed);
                PrepareSamples();
            }
        }

        public override IEnumerable<Vector2> GetSamples(Rectangle a_rect)
        {
            int top = (int)(a_rect.Top / m_cell_size) + 1;
            int bottom = (int)(a_rect.Bottom / m_cell_size) + 1;
            int left = (int)(a_rect.Left / m_cell_size) + 1;
            int right = (int)(a_rect.Right / m_cell_size) + 1;

            for (int x = Math.Max(0, left - 2); x <= Math.Min(m_grid_width - 1, right + 2); x++)
            {
                for (int y = Math.Max(0, top - 2); y <= Math.Min(m_grid_height - 1, bottom + 2); y++)
                {
                    if (m_grid[x, y] != -1)
                    {
                        Vector2 p = m_samples[m_grid[x, y]];

                        if (p.X < a_rect.Left)
                            continue;
                        if (p.X > a_rect.Right)
                            continue;
                        if (p.Y < a_rect.Top)
                            continue;
                        if (p.Y > a_rect.Bottom)
                            continue;

                        yield return p;
                    }
                }
            }
        }

        public override SamplerType SamplerType
        {
            get
            {
                return SamplerType.Poisson;
            }
        }
    }
}
